import os
import glob
from skimage import io
from nets.sesf_net import SESF_Fuse


def main(input_dir, output_dir):
    """
    Image Fusion
    :param input_dir: str, input dir with all images stores in one folder
    :param output_dir: str, output dir with all fused images
    :return:
    """
    sesf = SESF_Fuse("cse")
    images_name = sorted(list({item[:-6] for item in os.listdir(input_dir)}))
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    for image_name in images_name:
        print("Fusing {}".format(image_name))

        img1_name = os.path.join(input_dir, image_name + "_1.png")
        img2_name = os.path.join(input_dir, image_name + "_2.png")

        if not os.path.exists(img1_name):
            img1_name = os.path.join(input_dir, image_name + "_1.jpg")
            img2_name = os.path.join(input_dir, image_name + "_2.jpg")

        if not os.path.exists(img1_name):
            img1_name = os.path.join(input_dir, image_name + "_1.bmp")
            img2_name = os.path.join(input_dir, image_name + "_2.bmp")

        img1 = io.imread(img1_name)
        img2 = io.imread(img2_name)

        fused = sesf.fuse(img1, img2)
        io.imsave(os.path.join(output_dir, image_name + ".png"), fused)


def main_multi_images():       
    
    sesf = SESF_Fuse("cse")

    data_root = "/home/sean/Data/2020.01.15/test_connector2/VHX"   
    imgFiles = glob.glob(os.path.join( data_root + "/*.jpg")) 
    imgFiles.sort()

    for fileName in imgFiles:
        print("{}".format(fileName))
        
    
    fusedImg = sesf.fuseMultiImgs(imgFiles)
   
    io.imsave(os.path.join(data_root, "fused_color.png"), fusedImg)
    


if __name__ == "__main__":
    
    # test paired-image fusing
    '''
    input_dir = os.path.join(os.getcwd(), "data", "multi_focus")
    output_dir = os.path.join(os.getcwd(), "data", "result")
    main(input_dir, output_dir)
    '''

    # test multi-image fusing 
    main_multi_images()
